# $Id: ip.py 9227e81e919a 2012/08/27 15:49:57 pgurumur $
# Copyright (c) 2009 Prabhu Gurumurthy  <pgurumur@gmail.com>
# 
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

# XXX: Parts of this code, marked with %{ %} are under 
# Copyright 2007 Google Inc.
#  Licensed to PSF under a Contributor Agreement.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing
# permissions and limitations under the License.

__version__  = "$Revision: 9227e81e919a $"
__author__   = "$Author: pgurumur $"
__modified__ = "$Date: 2012-08-27 15:49:57Z $"

# import everything from netaddr 0.7.x
from netaddr import IPNetwork, IPAddress, AddrFormatError

try:
   from . import exception
except ImportError:
   import exception

# unified IPv4/IPv6 class for handling network and address type for both the
# families
class ip(object):
   # Class takes in one argument, this argument can be an IPv4/IPv6 subnet
   # or it can be an IPv4/IPv6 address
   def __init__(self, IPAddr, **Keywords):

      # Instantiate the exception class
      self._err = exception.IPError()

      # Main instance of either CIDR/IP class from netaddr
      self._ip = None

      # Whether given argument is a network or an address
      self._isNetwork = False

      if IPAddr is not None:
         if isinstance(IPAddr, (IPNetwork, ip)):
            self._ipAddr(str(IPAddr))
         elif isinstance(IPAddr, IPAddress):
            self._ip = IPAddr
         elif isinstance(IPAddr, int):
            if "version" in Keywords:
               ipaddr = None
               # This function takes in a integer and converts into IPaddress
               # if prefix is given, then make it into a network and 
               # instantiates the same class, raises exception.exception.IPError,
               # if there is a problem
               # Works for both IPv4 and IPv6
               if int(Keywords["version"]) == 4:
                  # For IPv4 addresses usually
                  octets = []
                  for _ in range(4):
                     octets.insert(0, str(IPAddr & 0xFF))
                     IPAddr >>= 8

                  ipaddr = ".".join(octets)
               elif int(Keywords["version"]) == 6:
                  """
                  # For IPv6 addresses usually
                  hexstr = "%32x" %IPAddr
                  hextets = []
                  for ix in range(0, 32, 4):
                     hextets.append("%x" %int(hexstr[ix:ix + 4], 16))

                  hextets = self._compress(hextets)
                  ipaddr = ":".join(hextets)
                  """

                  hextets = '0' * 32 + hex(IPAddr)[2:]
                  temp = ""
                  for ix in range(1, 33):
                     temp = hextets[-ix] + temp
                     if ix % 4 == 0:
                        temp = ':' + temp

                  ipaddr = temp[1:]
               else:
                  raise exception.IPError("unknown IP version")

               if "prefix" in Keywords:
                  if Keywords['prefix']:
                     ipaddr = "%s/%d" %(ipaddr, Keywords["prefix"])

               self._ipAddr(ipaddr)
            else:
               raise exception.IPError("no version defined!")
         elif isinstance(IPAddr, str):
            self._ipAddr(str(IPAddr))
         else:
            raise exception.IPError("unknown object instance: %s" %type(IPAddr))
      else:
         raise exception.IPError("no ip address/subnet defined!")

   def _ipAddr(self, Str):
      # Function that forms either IPNetwork or IPaddress instantiation
      # based on given string, used for constructor with following objects
      # IPNetwork
      # IPAddress
      # ip
      # String
      if isinstance(Str, str):
         # Check to see whether we have a forward slash, if we do, mostly
         # it is a network, we still will verify the prefixlen if the 
         # prefixlen is 32 or 128, then it is automatically converted into
         # IPAddress format
         iplen = len(Str.split("/"))

         try:
            if iplen == 2:
               # String with forward slash
               self._ip = IPNetwork(Str, implicit_prefix = True)

               prefix = self._ip.prefixlen
               # if the prefixlen is 32 or 128 it is an IPAddress
               if (prefix == 32) or (prefix == 128):
                  self._ip = IPAddress(Str.split("/")[0])
                  # The following if block is necessary, otherwise for a 
                  # /32 bit IPv6 address, it would be treated as an address
                  # instead of subnet
                  if (prefix == 32) and (self._ip.version != 4):
                     self._ip = IPNetwork(Str, implicit_prefix = True)
                     self._isNetwork = True
               else:
                  # Toggle to the network flag
                  self._isNetwork = True
            elif iplen == 1:
               self._ip = IPAddress(Str)
            else:
               raise exception.IPError("invalid IPv4/IPv6 address: %s" %Str)
         except ValueError as message:
            raise exception.IPError(message)
         except AddrFormatError as message:
            raise exception.IPError(message)

   def _compress(self, hextets):
      # %{
      # From ipaddr.py
      if hextets:
         """
         Compresses a list of hextets.
         Compresses a list of strings, replacing the longest continuous
         sequence of "0" in the list with "" and adding empty strings at
         the beginning or at the end of the string such that subsequently
         calling ":".join(hextets) will produce the compressed version of
         the IPv6 address.

         Args:
               hextets: A list of strings, the hextets to compress.

         Returns:
               A list of strings.

         """
         best_doublecolon_start = -1
         best_doublecolon_len = 0
         doublecolon_start = -1
         doublecolon_len = 0
         for index in range(len(hextets)):
            if hextets[index] == '0':
               doublecolon_len += 1
               if doublecolon_start == -1:
                  # Start of a sequence of zeros.
                  doublecolon_start = index
               if doublecolon_len > best_doublecolon_len:
                  # This is the longest sequence of zeros so far.
                  best_doublecolon_len = doublecolon_len
                  best_doublecolon_start = doublecolon_start
            else:
               doublecolon_len = 0
               doublecolon_start = -1

         if best_doublecolon_len > 1:
            best_doublecolon_end = (best_doublecolon_start +
                  best_doublecolon_len)

            # For zeros at the end of the address.
            if best_doublecolon_end == len(hextets):
               hextets += ['']

            hextets[best_doublecolon_start:best_doublecolon_end] = ['']
            # For zeros at the beginning of the address.
            if best_doublecolon_start == 0:
               hextets = [''] + hextets

      # %}
      return hextets

   def subnets(self, SubnetValue = 0):
      try:
         temp = int(SubnetValue)
      except ValueError as message:
         raise exception.IPError(message)
      else:
         if self._isNetwork:
            try:
               for items in list(self._ip.subnet(SubnetValue)):
                  yield ip(items)
            except exception.IPError as message:
               raise exception.IPError(message)

   def netmask(self, Wildcard = False):
      # Only IPv4 functions and the names are self explantory
      retval = None
      if self._isNetwork and (self._ip.version == 4):
         if Wildcard:
            retval = self._ip.hostmask
         else:
            retval = self._ip.netmask

      return retval

   def broadcast(self):
      # Only IPv4 functions and the names are self explantory
      retval = None
      if self._isNetwork and (self._ip.version == 4):
         retval = self._ip.broadcast

      return retval

   # Get the size of the network
   def size(self):
      retval = 1
      if self._isNetwork:
         retval = self._ip.size

      return retval

   # Binary values (in bits) of the IPv4/IPv6 address
   def binary(self):
      if not self._isNetwork:
         return self._ip.bits()

   # used for 'in' operand
   def __contains__(self, IPAddr):
      retval = False
      if self._isNetwork:
         temp = None
         try:
            temp = ip(IPAddr)
         except exception.IPError as message:
            raise exception.IPError(message)
         else:
            address = None
            if temp._isNetwork:
               address = IPNetwork(str(IPAddr))
            else:
               address = IPAddress(str(IPAddr))

            if address in self._ip:
               retval = True

      return retval

   # for int function
   def __int__(self):
      retval = None
      if self._isNetwork:
         retval = self._ip.value
      else:
         retval = int(self._ip)

      return retval

   # For list function
   def __iter__(self):
      if self._isNetwork:
         try:
            for items in list(self._ip):
               yield ip(items)
         except IndexError as message:
            raise exception.IPError(message)
      else:
         yield "%s" %self

   # for len function
   def __len__(self):
      retval = 0
      if self._isNetwork:
         retval = self._ip.prefixlen
      else:
         if self.version == 4:
            retval = 32
         elif self.version == 6:
            retval = 128

      return retval

   # for str function
   def __str__(self):
      # return str(self._ip)
      return str(self._ip)

   def __repr__(self):
      return repr(str(self).encode('ascii'))

   def __index__(self):
      return int(self)

   def __add__(self, obj):
      if isinstance(obj, str):
         temp = str(self._ip) + "%s" %obj
         self._ipAddr(temp)
      else:
         raise exception.IPError("invalid type ('%s') to add" %type(obj))

      return self._ip

   def __getitem__(self, Key):
      retval = None
      if isinstance(Key, int):
         try:
            if self.size() > 1:
               retval = str(self._ip[int(Key)])
            else:
               retval = str(self._ip)
         except ValueError as message:
            raise exception.IPError(message)
         except IndexError as message:
            raise exception.IPError(message)
      else:
         raise exception.IPError(
               "cannot get index value for non integer type key")

      return retval

   def __eq__(self, other):
      retval = False
      if isinstance(other, int) and int(self) == int(other):
         retval = True
      elif isinstance(other, str) and str(self) == str(other):
         retval = True

      return retval

   @property
   def reverse(self):
      retval = None
      if len(self) == 32:
         retval = self._ip.reverse_dns.split(".in-addr")[0]
      elif len(self) == 128:
         retval = self._ip.reverse_dns.split(".ip6")[0]
         
      return retval

   ismulticast = property(fget = lambda self: self._ip.is_multicast())
   isreserved = property(fget = lambda self: self._ip.is_reserved())
   version = property(fget = lambda self: self._ip.version)
   value = property(fget = lambda self: int(self))
   length = property(fget = lambda self: len(self))
   private = property(fget = lambda self: self._ip.is_private())
